/*
 * Copyright 2010 Google Inc.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.google.websocket;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Map;
import java.util.Random;
import java.util.Set;

/**
 * Implements the draft WebSocket handshake.
 */
public class DraftHandshake extends HttpCompatibleHandshake {

  private static final MessageDigest md5;

  static {
    MessageDigest md5temp = null;
    try {
      md5temp = MessageDigest.getInstance("MD5");
    } catch (NoSuchAlgorithmException e) {
    }
    md5 = md5temp;
  }

  /**
   * @return a v76-v00-style key value.
   */
  private static String generateWebSocketKey() {
    Random rng = new Random();
    int spaces = rng.nextInt(12) + 1;
    int max = Integer.MAX_VALUE / spaces;
    int number = rng.nextInt(max);
    int product = number * spaces;
    StringBuilder buf = new StringBuilder(String.valueOf(product));
    int randomChars = rng.nextInt(12) + 1;
    for (int i = 0; i < randomChars; ++i) {
      int pos = rng.nextInt(buf.length() + 1);
      int chIndex = rng.nextInt(84);
      char ch = (char) (chIndex >= 15 ? 0x003A + chIndex - 15
          : 0x0021 + chIndex);
      buf.insert(pos, ch);
    }
    for (int i = 0; i < spaces; ++i) {
      int pos = rng.nextInt(buf.length() - 1) + 1;
      buf.insert(pos, ' ');
    }
    return buf.toString();
  }

  public WebSocket clientHandshake(InputStream in, OutputStream out,
      Set<Extension> extensions, URI uri, String origin) throws IOException {
    String path = uri.getPath();
    if (path == null || path.length() == 0) {
      path = "/";
    }
    String query = uri.getQuery();
    if (query != null && query.length() > 0) {
      path += "?" + query;
    }
    StringBuilder req = new StringBuilder();
    req.append("GET ").append(path).append(" HTTP/1.1\r\n");
    req.append("Upgrade: WebSocket\r\n");
    req.append("Connection: Upgrade\r\n");
    String host = uri.getHost();
    if (uri.getPort() >= 0) {
      host += ":" + uri.getPort();
    }
    req.append("Host: ").append(host).append("\r\n");
    req.append("Origin: ").append(origin).append("\r\n");
    req.append("SEC-WebSocket-Key1: ").append(generateWebSocketKey()).append(
        "\r\n");
    req.append("SEC-WebSocket-Key2: ").append(generateWebSocketKey()).append(
        "\r\n");
    int draft = getDraftNumber();
    if (draft >= 0) {
      req.append("Sec-WebSocket-Draft: ").append(draft).append("\r\n");
    }
    req.append("\r\n");
    out.write(req.toString().getBytes(UTF8));
    out.flush();
    String line = Util.readUtf8Line(in);
    if (!"HTTP/1.1 101 WebSocket Protocol Handshake".equals(line)) {
      throw new IOException("Bad WebSocket response: " + line);
    }
    Map<String, String> headers = Util.readHeaderFields(in);
    if (!"Upgrade".equals(headers.get("connection"))
        || !"WebSocket".equals(headers.get("upgrade"))) {
      throw new IOException("Bad WebSocket response headers: " + headers);
    }
    byte[] clientNonce = new byte[16];
    new Random().nextBytes(clientNonce);
    out.write(clientNonce);
    byte[] serverNonce = new byte[16];
    Util.readFully(serverNonce, in);
    // TODO(jat): verify random bytes from server
    return createWebSocket(in, out, extensions);
  }

  @Override
  public boolean matches(String requestLine, Map<String, String> headers) {
    if (!"Upgrade".equalsIgnoreCase(headers.get("connection"))) {
      return false;
    }
    if (!"WebSocket".equalsIgnoreCase(headers.get("upgrade"))) {
      return false;
    }
    if (!headers.containsKey("host")) {
      return false;
    }
    if (!headers.containsKey("origin")) {
      return false;
    }
    int draft = getDraftNumber();
    if (draft > 0) {
      String draftValue = headers.get("sec-websocket-draft");
      if (draftValue == null || Integer.valueOf(draftValue) != draft) {
        return false;
      }
    } else {
      if (headers.containsKey("sec-websocket-draft")) {
        return false;
      }
    }
    return true;
  }

  @Override
  public WebSocket serverHandshake(String reqLine, Map<String, String> headers,
      InputStream in, OutputStream out, Set<Extension> extensions, int port)
      throws IOException {
    String[] get = reqLine.split(" ");
    if (get.length != 3 || !"GET".equals(get[0])) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    String resource = get[1];
    // TODO(jat): validate resource
    if (!"Upgrade".equalsIgnoreCase(headers.get("connection"))) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    if (!"WebSocket".equalsIgnoreCase(headers.get("upgrade"))) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    if (!headers.containsKey("host")) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    String host = headers.get("host");
    if (!headers.containsKey("origin")) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    String origin = headers.get("origin");
    String protocol = headers.get("sec-websocket-protocol");
    if (protocol != null) {
      throw new IOException("Unsupported subprotocol " + protocol);
    }
    String draft = headers.get("sec-websocket-draft");
    String key1 = headers.get("sec-websocket-key1");
    String key2 = headers.get("sec-websocket-key2");
    int key1number = extractKeyNumber(key1);
    int key2number = extractKeyNumber(key2);
    int key1spaces = extractSpaces(key1);
    int key2spaces = extractSpaces(key2);
    if (key1spaces == 0 || key2spaces == 0) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    int key1part = key1number / key1spaces;
    int key2part = key2number / key2spaces;
    if (key1part * key1spaces != key1number
        || key2part * key2spaces != key2number) {
      throw new IOException("Incorrect WebSocket handshake: k1='" + key1
          + "', k2='" + key2 + "', key1num=" + key1number + ", key1spaces="
          + key1spaces + ", key2num=" + key2number + ", key2spaces="
          + key2spaces);
    }
    out.write("HTTP/1.1 101 WebSocket Protocol Handshake\r\n".getBytes(UTF8));
    out.write("Connection: Upgrade\r\n".getBytes(UTF8));
    out.write("Upgrade: WebSocket\r\n".getBytes(UTF8));
    String line = "Sec-WebSocket-Location: " + host + ":" + port + resource
        + "\r\n";
    out.write(line.getBytes(UTF8));
    line = "Sec-WebSocket-Origin: " + origin + "\r\n";
    out.write(line.getBytes(UTF8));
    if (draft != null) {
      line = "Sec-WebSocket-Draft: " + draft + "\r\n";
      out.write(line.getBytes(UTF8));
    }
    out.write("\r\n".getBytes(UTF8));
    out.flush();
    byte[] clientNonce = new byte[16];
    Util.readFully(clientNonce, in);
    byte[] intBytes = new byte[4];
    intToBytes(key1part, intBytes);
    md5.update(intBytes);
    intToBytes(key2part, intBytes);
    md5.update(intBytes);
    md5.update(clientNonce);
    out.write(md5.digest());
    out.flush();
    return createWebSocket(in, out, extensions);
  }

  /**
   * @param in
   * @param out
   * @return a connected WebSocket instance
   * @throws IOException 
   */
  protected WebSocket createWebSocket(InputStream in, OutputStream out,
      Set<Extension> extensions) throws IOException {
    return new DraftWebSocket(in, out, extensions);
  }

  /**
   * @return draft number, or -1 if none present
   */
  protected int getDraftNumber() {
    return 2;
  }

  /**
   * @param key
   * @return the value, ignoring non-digits, of the key
   */
  private int extractKeyNumber(String key) {
    int val = 0;
    for (int i = 0; i < key.length(); ++i) {
      char ch = key.charAt(i);
      if (ch >= '0' && ch <= '9') {
        val = val * 10 + Character.digit(ch, 10);
      }
    }
    return val;
  }

  /**
   * @param key
   * @return count of spaces in the key
   */
  private int extractSpaces(String key) {
    int count = 0;
    for (int i = 0; i < key.length(); ++i) {
      if (key.charAt(i) == ' ') {
        count++;
      }
    }
    return count;
  }

  /**
   * @param val
   * @param intBytes
   */
  private void intToBytes(int val, byte[] intBytes) {
    intBytes[0] = (byte) ((val >> 24) & 255);
    intBytes[1] = (byte) ((val >> 16) & 255);
    intBytes[2] = (byte) ((val >> 8) & 255);
    intBytes[3] = (byte) (val & 255);
  }
}
